/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunriset is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Implementation of grain species and size distribution classes.

// $Id$

#include "config.h"
#include "grain.h"
#include "CCfits/CCfits" 
#include "fits-utilities.h"
#include "blitz-fits.h"
#include "interpolatort.h"
#include "constants.h"
#include "threadlocal.h"
#include "vecops.h"
#include <cassert>

#ifdef WITH_CUDA
#include "cuda_grain_temp.h"
#endif

using namespace mcrx;
using namespace blitz;
using namespace std;
using namespace CCfits;
  

mcrx::grain_data::grain_data (const std::string& file_name, T_float minsize, 
			      T_float maxsize)
{
  load_file (file_name, minsize, maxsize); 

  asizes_.reference(reference_vector(sizes_));
  delta_size_.reference(delta_quantity(asizes_));

  if(!(all(sigma_>0)&&all(esigma_>0)&&all(asizes_>0))) {
    cerr << "Zero cross section or size encountered in grain data file: " 
	 << file_name << endl;
    throw 1;
  }

}


void mcrx::grain_data::load_file (const std::string& file_name, 
				  T_float minsize, T_float maxsize)
{
  cout << "Reading grain cross-sections from file "<< file_name << endl;
  FITS file (file_name, Read);

  read (open_HDU(file,"SIGMA_ABS"), sigma_);
  sigma_.reverseSelf(secondDim);
  {
    // seems to be a bug in blitz that makes reversed dimensions
    // iffy. copy. and make it a contiguous array since the cuda code
    // can't handle the padded one
    array_2 temp(sigma_.shape(), blitz::contiguousArray);
    temp=sigma_;
    sigma_.reference(temp);
  }
  read (open_HDU(file,"SIGMA_SCA"), sigma_sca_);
  sigma_sca_.reverseSelf(secondDim);
  read (open_HDU(file,"G"), g_);
  g_.reverseSelf(secondDim);

  ASSERT_ALL(sigma_>=0);
  ASSERT_ALL(sigma_sca_>=0);
  ASSERT_ALL(g_>=-1);
  ASSERT_ALL(g_<=1);

  ExtHDU& hdu = open_HDU (file, "AXES");
  hdu.column("size").read(sizes_, 1, sigma_.extent(firstDim ));
  hdu.column("wavelength").read(lambda_, 1, sigma_.extent(secondDim));
  reverse(lambda_.begin(), lambda_.end()); 

  // only select sizes we want
  const int min= lower_bound (sizes_.begin(), sizes_.end(),
			      minsize) - sizes_.begin();
  const int max= lower_bound (sizes_.begin(), sizes_.end(),
			      maxsize) - sizes_.begin();
  // note that the Range() and the iterator range works differently wrt end
  sigma_.reference(sigma_(Range(min,max-1), Range::all()));
  // esigma just references sigma until necessary
  esigma_.reference(sigma_);
  sigma_sca_.reference(sigma_sca_(Range(min,max-1), Range::all()));
  g_.reference(g_(Range(min,max-1), Range::all()));
  sizes_.assign(&sizes_[min], &sizes_[max]);
  assert(sizes_.size()==sigma_.extent(firstDim));
  cout << "  Selecting size range " << sizes_.front()
       << " - " << sizes_.back() << endl;

  // we assume that the grain file is written with consistent units
  units_["length"] = hdu.column("size").unit();
  units_["wavelength"] = hdu.column("wavelength").unit();
  // the unit of luminosity calculated for SED is W
  units_["luminosity"] = "W";

  // since this is called from the constructor, we can't make the
  // virtual call here, because the derived class has not been
  // constructed yet. NOTE that this means the derived class is
  // responsible for calling its own setup in its own constructor.
  grain_data::setup();
}

void mcrx::grain_data::resample(const array_1& new_lambda,
				const array_1& new_xelambda)
{
  // for SOME reason, we can't redefine an array passed by value as
  // one of the parameters, so we have to clone it.
  array_1 new_elambda(new_xelambda);

  if((new_lambda(0)<alambda_(0)) ||
     (new_lambda(new_lambda.size()-1)>alambda_(alambda_.size()-1))) {
    cerr << "Error: Trying to resample grain cross sections outside defined wavelengths:\n" << alambda_(0) << " - " << alambda_(alambda_.size()-1) << "\n" << new_lambda(0) << " - " << new_lambda(new_lambda.size()-1) << endl;
    throw false;
  }

  // for dumping out what the cross sections actually look like
  const bool dump_sigmas=false;
  if(dump_sigmas) {
    ofstream crap1("sigma.txt");
    ofstream crap2("sigma_s.txt");
    ofstream crap3("g.txt");
    for(int i=0; i <sigma_.extent(secondDim); ++i) {    
      crap1 << alambda_(i) << '\t';
      crap2 << alambda_(i) << '\t';
      crap3 << alambda_(i) << '\t';
      for(int j=0; j <sigma_.extent(firstDim); ++j) {    
	crap1 << sigma_(j,i) << '\t';
	crap2 << sigma_sca_(j,i) << '\t';
	crap3 << g_(j,i) << '\t';
      }
      crap1 << '\n';
      crap2 << '\n';
      crap3 << '\n';
    }
  }
  
  // gotta be careful about which arrays use which lambda here. The
  // sigma_s and g arrays use the *emission* grid, because they only
  // pertain to scattering. (Seems counterintuitive but "emission"
  // really means the emission and subsequent transfer of radiation
  // and "absorption" the calculation of heating.)
  bool sep_elambda=true;
  if( (new_elambda.size()==0) || 
      ( (new_elambda.size()==new_lambda.size()) && 
	all(new_elambda==new_lambda))) {
    new_elambda.reference(new_lambda);
    sep_elambda=false;
  }

  bool alambda_unchanged = 
    (new_lambda.size()==alambda().size()) && 
    all(new_lambda==alambda());
  bool elambda_unchanged =
    (new_elambda.size()==elambda().size()) && 
    all(new_elambda==elambda());
  if(alambda_unchanged && elambda_unchanged)
    return;

  // check if we are upsampling or downsampling. If upsampling, we use
  // mcrx::resample with a logarithmic resampling, because that works
  // better. For downsampling, we use mcrx::subsample, which preserves
  // the integral. Since the emission grid never smaller than the
  // absorption one, we use the emission grid to determine this
  const bool upsample = 
    (log10(alambda_(alambda_.size()-1))-log10(alambda_(0)))/(alambda_.size()-1) >
    ((log10(new_elambda(new_elambda.size()-1))-log10(new_elambda(0)))/
     (new_elambda.size()-1));
  if (upsample) 
    cout << "Upsampling grain opacities onto wavelength grid\n";
  else
    cout << "Downsampling grain opacities onto wavelength grid\n";

  // check if we have separate emission wavelengths. if we do we
  // resample the opacity onto those wavelengths, but to ensure
  // integral is conserved, we subsample the absorption opacities.
  array_2 new_esigma(sigma_.extent(firstDim), new_elambda.size(), 
		     contiguousArray);
  array_2 new_sigma(sigma_.extent(firstDim),
		    new_lambda.size(), contiguousArray);
  array_2 new_sigma_s(sigma_.extent(firstDim),
		      new_elambda.size());
  array_2 new_g(new_sigma_s.shape());

  if (!elambda_unchanged) {
    cout << "\tResampling emission quantities\n";
    if (upsample)
      for(int i=0; i <sigma_.extent(firstDim); ++i) {    
	new_esigma(i,Range::all()) = mcrx::resample(sigma_(i,Range::all()),
						    elambda_, new_elambda, 
						    false, true);
	new_sigma_s(i,Range::all()) = mcrx::resample(sigma_sca_(i,Range::all()),
						     elambda_, new_elambda, 
						     false, true);
      }
    else
      for(int i=0; i <sigma_.extent(firstDim); ++i) {    
	new_esigma(i,Range::all()) = mcrx::subsample(sigma_(i,Range::all()),
						     elambda_, new_elambda);
	new_sigma_s(i,Range::all()) = mcrx::subsample(sigma_sca_(i,Range::all()),
						      elambda_, new_elambda);
      }

    for(int i=0; i <sigma_.extent(firstDim); ++i) {
      // g can be negative so can't use log sampling anyway
      new_g(i,Range::all()) = mcrx::subsample(g_(i,Range::all()),
					      elambda_, new_elambda);
    }

    ASSERT_ALL(new_esigma>=0);
    ASSERT_ALL(new_sigma_s>=0);
    ASSERT_ALL(new_g>=-1);
    ASSERT_ALL(new_g<=1);

    elambda_.reference(new_elambda);
    esigma_.reference(new_esigma);
    sigma_sca_.reference(new_sigma_s);
    g_.reference(new_g );
  }

  if(!alambda_unchanged || !elambda_unchanged) {
    cout << "\tResampling absorption quantities\n";
    // to make sure these have consistent cross sections, we always
    // *subsample* the absorption opacities from the emission ones.
    if (sep_elambda) {
      for(int i=0; i <sigma_.extent(firstDim); ++i)
	new_sigma(i,Range::all()) = mcrx::subsample(esigma_(i,Range::all()),
						    new_elambda, new_lambda);
      sigma_.reference(new_sigma);
    }
    else
      sigma_.reference(esigma_);
    ASSERT_ALL(sigma_>=0);
  }    

  lambda_.assign(new_lambda.begin(), new_lambda.end());

  if(dump_sigmas) {
    ofstream crap1("sigma2e.txt");
    ofstream crap2("sigma_s2.txt");
    ofstream crap3("g2.txt");
    for(int i=0; i <esigma_.extent(secondDim); ++i) {    
      crap1 << elambda_(i) << '\t';
      crap2 << elambda_(i) << '\t';
      crap3 << elambda_(i) << '\t';
      for(int j=0; j <esigma_.extent(firstDim); ++j) {    
	crap1 << esigma_(j,i) << '\t';
	crap2 << sigma_sca_(j,i) << '\t';
	crap3 << g_(j,i) << '\t';
      }
      crap1 << '\n';
      crap2 << '\n';
      crap3 << '\n';
    }
    ofstream crap1a("sigma2a.txt");
    for(int i=0; i <sigma_.extent(secondDim); ++i) {    
      crap1a << alambda_(i) << '\t';
      for(int j=0; j <sigma_.extent(firstDim); ++j) {    
	crap1a << sigma_(j,i) << '\t';
      }
      crap1a << '\n';
    }

  }

  // and calculate the derived internal variables. here we make a
  // virtual call because the derived classes may need to update too.
  setup();
}

void mcrx::grain_data::setup()
{
  // This function must be safe to call multiple times because it's
  // called both from grain_data constructor and the derived setup.
  alambda_.reference(reference_vector(lambda_));
  if(elambda_.size()==0)
    elambda_.reference(alambda_);
  invelambda_.reference(make_thread_local_copy (1./elambda_));
  //delta_lambda_.reference(delta_quantity (alambda_));
}



/** Central wavelengths of PAH transitions in cm^-1. */
T_float mcrx::Brent_PAH_grain::x0_[]= {
  3040.3, 1897.0, 1754.0, 1608.5, 1608.5,
  1593.9, 1490.0, 1400.0, 1313.0, 1270.0,
  1200.0, 1163.1, 998.0 , 940.0 , 890.0 ,
  883.0 , 836.0 , 813.0 , 800.0 , 788.0 ,
  737.0 , 702.0 , 670.0 , 635.0 , 607.0 ,
  588.0 , 576.0 , 571.0 , 562.0 , 530.0};

/** Strengths of PAH transitions (in arbitrary units). */
T_float mcrx::Brent_PAH_grain::f0_[] = {
  1.006E-04, 1.000E-04, 1.000E-04, 3.420E-04, 4.600E-04,
  2.140E-04, 5.000E-05, 3.420E-04, 1.200E-03, 1.280E-03,
  3.200E-04, 8.900E-04, 1.630E-04, 1.600E-04, 1.700E-03,
  1.800E-03, 4.280E-04, 2.000E-04, 5.300E-04, 1.200E-03,
  3.600E-04, 2.570E-04, 8.550E-05, 1.710E-04, 7.800E-04,
  8.300E-04, 4.280E-04, 2.140E-04, 8.550E-05, 1.450E-04};

/** Widths of PAH transitions in cm^-1. */
T_float mcrx::Brent_PAH_grain::sigma_[] = {
  22.4, 40.0, 40.0, 37.8, 14.4,
  34.9, 30.0, 100.0, 28.0, 35.0,
  30.0, 27.0, 129.1, 13.0, 4.0,
  14.1, 14.1, 15.0, 70.7, 13.0,
  18.2, 12.9, 18.3, 18.3, 7.5,
  5.8, 4.0, 20.0, 3.8, 7.0};

void mcrx::Brent_PAH_grain::setup()
{
  grain_data::setup();

  // Now calculate the PAH template as a sum of Lorenzians. To make
  // sure we get a correct result if the wavelength resolution is low
  // we generate a supersampled version first and then subsample.
  const T_float dr=
    log10(elambda()(elambda().size()-1))-log10(elambda()(0));
  const int n_templambda= (elambda().size()<100) ? 1000 : 10*elambda().size();
  array_1 templambda(n_templambda);
  templambda = pow(10, tensor::i*dr/(n_templambda-1)+log10(elambda()(0)));
  // make sure the end points are *identical* regardless of truncation
  templambda(0) = elambda()(0);
  templambda(n_templambda-1) = elambda()(elambda().size()-1);

  template_emission_.resize(templambda.shape());
  template_emission_= 0;
  const int n = 30;
  const T_float lambda_conv = 
    units::convert(units().get("wavelength"), "cm");
  for (int i = 0; i < n; ++i) {
    template_emission_ += 
      f0_[i]/(1.0 +
	      (1./(lambda_conv*templambda)-x0_[i])*
	      (1./(lambda_conv*templambda)-x0_[i])/
	      (sigma_[i]*sigma_[i])
	      );
  }
  // new resample it to desired resolution
  template_emission_.reference(subsample(template_emission_, templambda, elambda()));

  // Brent's template is in F_nu, we want F_lambda, so we need to
  // multiply by 1/lambda^2. To get units right, we would also need to
  // multiply by c0, but since the template is normalized anyway,
  // that's unnecessary.

  template_emission_ *= invelambda()*invelambda();

  // Normalize to unit luminosity
  template_emission_ /= integrate_quantity (template_emission_, elambda ());

  threadLocal_warn(template_emission_, true);

  ofstream crap("pahtemplate");
  for(int i=0;i<alambda().size();++i)
    crap << alambda()(i) << '\t' << template_emission_(i) << endl;
}

